r"""General solver module for all algorithms implemented in PyHCL.

The general form of the integral conservation law solved by PyHCL can
be written as:

.. math::

    \frac{d\tilde{U}}{dt} =
    -\frac{1}{|V|}\sum_{s=1}^{N}\int_{A_s}^{A_s +1} \mathcal{H}\cdot
    n_s dA_s

In this formulation, the solution :math:`U` is assumed to be cell
averaged values residing at the centroid of each cell as shown in the
figure:

.. figure:: ../../../docs/source/images/triangle.png
    :align: center

This results in a cell-centered finite volume scheme which is uniquely
determined by the choice of numerical flux at the interface of two cell.

In 1D, for the gird structure as shown in the figure,

.. figure:: ../../../docs/source/images/1dcells.png
    :align: center

this formulation reduces to the familiar conservative formula:

.. math::

     Q_{i}^{n+1} = Q_{i}^n - \frac{\Delta t}{\Delta x} \left [
     F_{i+\frac{1}{2}} - F_{i - \frac{1}{2}} \right ]

Where :math:`F_{i + \frac{1}{2}}` is the numerical flux for the
interface at :math:`x_{i + \frac{1}{2}}`

Thus, given a grid or mesh data structure, we iterate over the faces
of the mesh and store a flux for that face based on the cell values in
nearby cells.

The :class:`Solver` class is the base class for all the solvers in
PyHCL. It provides a unified interface with provisions to set general
solver related constants like the CFL condition and I/O related
operations like setting the output directory, output filename and
output frequency.

"""

from optparse import OptionParser, OptionGroup, Option
from os.path import basename, splitext, abspath, join
import sys
import numpy

from utils import mkdir, PBar

from hogs.solvers.bc.boundary_conditions import BoundaryCondition

class Solver1D(object):
    r"""Base class for all 1D solvers.

    .. attribute:: nvar

     The number of conserved variables for the system. For example,
     the one dimensional Euler equations have :math:`m = 3`

    .. attribute:: tf

     Final time for the simulation. This can be overwritten by the
     command line arguments.

     .. attribute:: fname

      File name for the output.

    """
    def __init__(self, nvar, tf=1, pfreq=100, grid=None, fname=None):
        self.tf = tf
        self.nvar = nvar

        self.grid = grid

        # output path directory
        self.path = None

        # set the print frequency
        self.pfreq = pfreq

        # setup the default boundary condition
        self.bc = BoundaryCondition()

        # set the CFL number
        self.cfl = 0.6

        # set the primitive variable function
        self.primitive_variable_function = None

        # set the default time step function
        self.time_step_function = None

        # filename for the output
        if fname is None:
            fname = sys.argv[0].split('.')[0]

        self.fname = fname

        # the remaining are the arguments
        self.args = sys.argv[1:]

        # setup the option parser
        self._setup_optparse()

    def solve(self):
        t = 0.0
        count = 0

        bar = PBar(maxval=101)
        _tf = 1.0/self.tf

        while t < self.tf:

            # update primitive variables
            self.update_primitive_variables()

            # compute the time step
            dt = self.compute_dt()
            #print "Using dt = %g"%(dt)

            # I/O routines
            if count % self.pfreq == 0:
                self.save(count, dt, t)

            # apply boundary conditions
            self.apply_boundary_conditions()

            # update
            self.update(dt)

            # update the progressbar
            percentage_done = int ( t/self.tf * 100 )
            bar.update( percentage_done )

            count += 1
            t += dt

        bar.finish()

        self.save(count, dt, t)

    def update(self, dt):
        """Advance the solution using the standard conservative formula.

        ..math::

        q_{i}^{n+1} = q_{i}^n - frac{\Delta t}{\Delta x}
        (F_{i+\frac{1}{2}} - F_{i - \frac{1}{2}}

        """

        grid = self.grid
        k = dt/grid.dx

        # compute the fluxes.
        self.flux_function.compute_flux(grid, dt)

        # Update the cell values
        for cell in grid.cells[grid.nb:-grid.nb]:
            fl, fr = cell.get_fluxes()

            q = cell.q
            q[:] = q - k * (fr - fl)            

    def _setup_optparse(self):
        usage = """
        python %prog [options] """
        self.parser = parser = OptionParser(usage)

        # Add some default options

        parser.add_option("--directory", action="store", type="string",
                         dest="outdir", default=self.fname+'_output',
                         help="""The output directory""")

        parser.add_option("--fname", action="store", type="string",
                         dest="fname", default=self.fname,
                         help="Output file name")

        parser.add_option("--pfreq", action="store", type="int",
                         dest="pfreq", default=self.pfreq,
                         help="Output print frequency")

        numerics = OptionGroup(parser, "Numerics", "Solver options for numerics")

        numerics.add_option("--tf", action="store", type="float",
                            dest="tf", default=self.tf,
                            help="Final time")

        numerics.add_option("--cfl", action="store", type="float",
                            dest="cfl", default=0.6,
                            help="CFL number to use. Default 0.6")

        parser.add_option_group(numerics)

    def _process_command_line(self):
        (options, args) = self.parser.parse_args(self.args)
        self.options = options
        
        #save the path where we want to dump output
        self.path = abspath(options.outdir)
        mkdir(self.path)

    def add_option(self, opt):
        """ Add an Option/OptionGroup or their list to OptionParser """
        if isinstance(opt, OptionGroup):
            self.parser.add_option_group(opt)
        elif isinstance(opt, Option):
            self.parser.add_option(opt)
        else:
            # assume a list of Option/OptionGroup
            for o in opt:
                self.add_option(o)        

    def setup(self):
        """Perform all solver related setup.

        This function should be called just before calling the
        :function:`solve` method of the solver.

        We sequentially query all objects for solver specific options,
        before proceeding to process them.

        """
        solver_opts = self.get_options(self.parser)
        if solver_opts is not None:
            self.add_option(solver_opts)

        # process all command line options
        self._process_command_line()

        # set the individual options
        options = self.options

        # perform any solver specific setup
        self.set_options(options.__dict__)

        # set the cfl number
        self.set_cfl(options.cfl)

        # set the final time
        self.set_final_time(options.tf)

        # set the print frequency
        self.set_print_frequency(options.pfreq)

        # set the output filename
        self.set_output_fname(options.fname)

    def save(self, count, dt, t):

        if type(count) == str:
            fnamebase = self.fname + "_%s"%count
        else:
            fnamebase = self.fname + "_%03d"%count

        fname = join( self.path, fnamebase )

        # get the face coordinates and cell centers
        x = self.grid.x
        xc = 0.5 * (x[1:] + x[:-1])

        numpy.savez(fname,  t=t, dt=dt, x=x, xc=xc,
                    q=self.grid.q, w=self.grid.w)        

    def apply_boundary_conditions(self):
        """Apply the boundary conditions"""
        self.bc.apply_bc( self.grid )
        #self.bc_left.apply_bc(self.grid)
        #self.bc_right.apply_bc(self.grid)

    def update_primitive_variables(self):
        """Update the primitive variables"""
        if self.primitive_variable_function:
            self.primitive_variable_function.update_primitive_variables()

    def compute_dt(self):
        """Compute the time step"""
        smax = self.time_step_function.get_max_speed()
        grid = self.grid

        x = grid.x
        dx = min( x[1:] - x[:-1] )

        return self.cfl * dx/smax

    def set_print_frequency(self, pfreq):
        """Set the output print frequency"""
        self.pfreq = pfreq

    def set_output_fname(self, fname):
        """Set the outpt file name"""
        self.fname = fname

    def set_final_time(self, tf):
        """Set the final time"""
        self.tf = tf

    def set_cfl(self, cfl):
        """Set the CFL number to use"""
        self.cfl = cfl

    def initialize_grid(self, xlow=0.0, xhigh=1.0, dx=0.01, nb=1):
        self.grid.initialize(xlow, xhigh, dx, self.nvar, nb)

    def get_options(self, parser):
        """Setup options for the EELagrangianSolver"""
        opt = self.flux_function.get_options(parser)
        return opt

    def set_options(self, options):
        """Set the options for the EELagrangianSolver"""
        self.flux_function.set_options(options)
